import retro 
import numpy as np
import cv2
import neat
import pickle
from random import random
from time import sleep

# Play this retro game at this level.
env = retro.make('HangOn-Sms')

imgarray = []

# 加载配置文件
config = neat.Config(neat.DefaultGenome, neat.DefaultReproduction,
                     neat.DefaultSpeciesSet, neat.DefaultStagnation,
                    'D:/2024/hang-on/config-feedforward-2')


checkpoint_file = 'D:/2024/毕设/优秀训练结果/neat-checkpoint-873'    #n为第n次训练结果
p = neat.Population(config)
p = neat.Checkpointer.restore_checkpoint(checkpoint_file)

# 创建种群
# p = neat.Population(config)


# # 添加报告器
# p.add_reporter(neat.StdOutReporter(True))
# stats = neat.StatisticsReporter()
# p.add_reporter(stats)
# p.add_reporter(neat.Checkpointer(10))

def play_game( genomes, config):
    for genome_id, genome in genomes:
        score = 0
        ob = env.reset()
        iny, inx, inc = env.observation_space.shape

        # 将观测空间的大小缩小一倍
        inx = int(inx/8)
        iny = int(iny/8)

        done = False
        net = neat.nn.recurrent.RecurrentNetwork.create(genome, config)

        while not done:
            env.render()
            ob = cv2.resize(ob, (inx, iny))
            ob = cv2.cvtColor(ob, cv2.COLOR_BGR2GRAY)
            ob = np.reshape(ob, (inx, iny))
            imgarray = ob.flatten()
            # 使用神经网络输出动作
            nnOutput = net.activate(imgarray)
            nnOutput[0] = 0 #刹车   
            nnOutput[8] = 1  #油门
            ob, reward, done, info = env.step(nnOutput)
            score += reward
            frame = env.render(mode='rgb_array')
            sleep(0.001)
        print('Done. Score: {}'.format(score))

# 运行种群
# winner = p.run(eval_genomes)
play = p.run(play_game)

env.close()